996. Number of Squareful Arrays.go

package leetcode

import (
"math"
"sort"
)

funcnumSquarefulPerms(A []int) int {
iflen(A) ==0 {
return0
 }
used, p, res:=make([]bool, len(A)), []int{}, [][]int{}
sort.Ints(A) // 这里是去重的关键逻辑
generatePermutation996(A, 0, p, &res, &used)
returnlen(res)
}

funcgeneratePermutation996(nums []int, indexint, p []int, res*[][]int, used*[]bool) {
ifindex==len(nums) {
checkSquareful:=true
fori:=0; i<len(p)-1; i++ {
if!checkSquare(p[i] +p[i+1]) {
checkSquareful=false
break
 }
 }
ifcheckSquareful {
temp:=make([]int, len(p))
copy(temp, p)
*res=append(*res, temp)
 }
return
 }
fori:=0; i<len(nums); i++ {
if!(*used)[i] {
ifi>0&&nums[i] ==nums[i-1] &&!(*used)[i-1] { // 这里是去重的关键逻辑
continue
 }
iflen(p) >0&&!checkSquare(nums[i]+p[len(p)-1]) { // 关键的剪枝条件
continue
 }
 (*used)[i] =true
p=append(p, nums[i])
generatePermutation996(nums, index+1, p, res, used)
p=p[:len(p)-1]
 (*used)[i] =false
 }
 }
return
}

funccheckSquare(numint) bool {
tmp:=math.Sqrt(float64(num))
ifint(tmp)*int(tmp) ==num {
returntrue
 }
returnfalse
}